# 这里就是直接调用IK，传参是mocap点的位置和当前关节空间位置
# 然后把结果在mujoco里面显示出来

import mujoco
import mujoco.viewer as viewer
import numpy as np
from os.path import dirname, join
import time

# 导入自定义函数和CANFD SDK动态链接库>>>
from os.path import dirname, join
import sys
import platform
arch = platform.machine()
sys.path.append(join(dirname(__file__), 'src/lysdemo_'+arch))
sys.path.append(join(dirname(__file__), "src"))
sys.path.append(dirname(__file__))
from src import *

pos_sensor_names = [
        "Motor_Shoulder_pos", 
        "Motor_BigArm1_pos",
        "Motor_SmallArm1_pos",
        "Motor_SmallArm2_pos",
        "Motor_Hand_pos",
        "Motor_Finger_pos"
        ]

vel_sensor_names = [
        "Motor_Shoulder_vel", 
        "Motor_BigArm1_vel",
        "Motor_SmallArm1_vel",
        "Motor_SmallArm2_vel",
        "Motor_Hand_vel",
        "Motor_Finger_vel"
        ]

# mujoco加载模型
model_dof6 = mujoco.MjModel.from_xml_path(join(dirname(__file__),'scene_6DOF_new.xml')) 
data_dof6 = mujoco.MjData(model_dof6)
mocap_id_dof6 = model_dof6.body("target").mocapid[0]



with viewer.launch_passive(model_dof6,data_dof6) as viewer:
    # Initialize the camera view to that of the free camera.
    mujoco.mjv_defaultFreeCamera(model_dof6, viewer.cam)
    viewer.opt.frame = mujoco.mjtFrame.mjFRAME_SITE

    # 初始化位置
    data_dof6.qpos = np.array([1,1,1,1.5,-1,-0.5])
    data_dof6.ctrl = np.array([1,1,1,1.5,-1,-0.5])
    mujoco.mj_step(model_dof6, data_dof6)
    mujoco.mj_forward(model_dof6, data_dof6)

    while viewer.is_running():
        q_real_float = np.array([data_dof6.sensor(pos_sensor_name).data[0] for pos_sensor_name in pos_sensor_names])
        qv_real_float = np.array([data_dof6.sensor(vel_sensor_name).data[0] for vel_sensor_name in vel_sensor_names])

        mocap_pos = data_dof6.mocap_pos[mocap_id_dof6]

        mocap_quat = data_dof6.mocap_quat[mocap_id_dof6]

        start = time.time()

        data_dof6.ctrl = ika_admittance(q_real_float, mocap_pos, mocap_quat)# 展示ik位置
        print(f"mocap pos: {mocap_pos}")
        print(f"mocap quat: {mocap_quat}")
        print(f"ik + admittance time: {time.time()-start}")

        data_dof6.qpos = np.array([0.5,0.5,0.5,0.5,0.5,0.5])
        mujoco.mj_step(model_dof6, data_dof6)
        mujoco.mj_forward(model_dof6, data_dof6)

        viewer.sync()
        time.sleep(model_dof6.opt.timestep)